#include <linux/init.h>
#include <linux/module.h>
#include <linux/spi/spi.h>
#include <linux/mtd/mtd.h>
#include <linux/string.h>
#include <linux/gpio.h>
#include <linux/of.h>
#include <linux/of_gpio.h>
#include <linux/timer.h>
#include <linux/delay.h>
#include <linux/vmalloc.h>
#include "NM25Q128EVB.h"

struct nm25q128evb {
    struct spi_device *spi;
    struct mtd_info mtd;
    u16 deviceid;
    int power_gpio;
};

static u16 flash_read_deviceid(struct spi_device *spi)
{
    u16 device_id = 0;
    u8 cmd[4] = {FLASH_ManufactDeviceID, 0x0, 0x0, 0x0};
    u8 rx[2];

    spi_write_then_read(spi, cmd, sizeof(cmd), rx, sizeof(rx));
    device_id = rx[0] << 8 | rx[1];

    return device_id;
}

static int flash_set_address(u64 address, u8 *cmd, u32 len)
{
    if (len == 4) {
        cmd[0] = (u8)(address >> 24);
        cmd[1] = (u8)(address >> 16);
        cmd[2] = (u8)(address >> 8);
        cmd[3] = (u8)address;
    } else {
        cmd[0] = (u8)(address >> 16);
        cmd[1] = (u8)(address >> 8);
        cmd[2] = (u8)address;
    }
    return 0;
}

static void flash_write_cmd(struct spi_device *spi, u8 command)
{
    int ret;
    u8 reg[] = {command};
    ret = spi_write(spi, reg, sizeof(reg));
    if (ret) {
        dev_err(&spi->dev, "write cmd %u failed.ret=%d\n", command, ret);
    }
}

static void flash_write_sr(struct spi_device *spi, u8 command, u8 data)
{
    int ret;
    u8 reg[] = {command, data};
    ret = spi_write(spi, reg, sizeof(reg));
    if (ret) {
        dev_err(&spi->dev, "write status register %u failed.ret=%d\n", command, ret);
    }
}

static u8 flash_read_sr(struct spi_device *spi, u8 command)
{
    u8 reg[1] = {command};
    u8 data[1];
    int ret;
    ret = spi_write_then_read(spi, reg, sizeof(reg), data, sizeof(data));
    if (ret) {
        dev_err(&spi->dev, "read status register %u failed.ret=%d\n", command, ret);
    }
    return data[0];
}

static void flash_set_word(struct spi_device *spi)
{
    u8 temp;
    temp = flash_read_sr(spi, FLASH_ReadStatusReg3);
    if ((temp & 0x01) == 0) {
        flash_write_cmd(spi, FLASH_WriteEnable);
        temp |= 1 << 1;
        flash_write_sr(spi, FLASH_WriteStatusReg3, temp);
        flash_write_cmd(spi, FLASH_Enable4ByteAddr);
    }
}

static void rk3588_mtd_sync(struct mtd_info *mtd)
{
    struct spi_device *spi = mtd->priv;
    while ((flash_read_sr(spi, FLASH_ReadStatusReg1) & 0x01) == 0x01);
}

static int rk3588_mtd_erase(struct mtd_info *mtd, struct erase_info *instr)
{
    struct spi_device *spi = (struct spi_device*)mtd->priv;
    struct nm25q128evb *p = spi_get_drvdata(spi);
    u64 i;
    u32 sector_erase_size = p->deviceid == W25Q256 ? 4 : 3;
    u8 *sector_erase;
    u64 start_addr = (instr->addr / mtd->erasesize) * mtd->erasesize;
    u64 actual_len = instr->addr - start_addr + instr->len;
    int ret = 0;

    sector_erase = kmalloc(sizeof(u8) * (sector_erase_size + 1), GFP_KERNEL);
    if (!sector_erase) {
        dev_err(&spi->dev, "[%s]allocate sector_erase memory failed\n", __func__);
        ret = -ENOMEM;
        goto ERASE_SECTOR_ALLOCATE_FAILED;
    }

    flash_write_cmd(spi, FLASH_WriteEnable);
    for (i = start_addr; i < start_addr + actual_len; i += mtd->erasesize) {
        sector_erase[0] = FLASH_SectorErase;
        flash_set_address(i, &sector_erase[1], sector_erase_size);
        ret = spi_write(spi, sector_erase, sector_erase_size + 1);
        if (ret) {
            dev_err(&spi->dev, "erase sector %llu failed.ret=%d\n", i, ret);
            goto ERASE_FAILED;
        }
        rk3588_mtd_sync(mtd);
    }

ERASE_FAILED:
    kfree(sector_erase);
ERASE_SECTOR_ALLOCATE_FAILED:
    return ret;
}

static int rk3588_mtd_write(struct mtd_info *mtd, loff_t to, size_t len,
                            size_t *retlen, const u_char *buf)
{
    struct spi_device *spi = (struct spi_device*)mtd->priv;
    struct nm25q128evb *p = spi_get_drvdata(spi);
    u64 i;
    u8 *buff = (u8*)buf;
    u32 sector_write_size = p->deviceid == W25Q256 ? 4 : 3;
    u8 *sector_write;
    u64 start_addr = (to / mtd->writesize) * mtd->writesize;
    u64 actual_len = to - start_addr + len;
    u64 rest_len;
    int ret = 0;
    struct spi_message m;
    struct spi_transfer t[2];
    memset(t, 0, sizeof(struct spi_transfer) * 2);

    sector_write = kmalloc(sizeof(u8) * (sector_write_size + 1), GFP_KERNEL);
    if (!sector_write) {
        dev_err(&spi->dev, "[%s]allocate sector_write memory failed\n", __func__);
        ret = -ENOMEM;
        goto WRITE_SECTOR_ALLOCATE_FAILED;
    }

    flash_write_cmd(spi, FLASH_WriteEnable);
    for (i = start_addr, rest_len = actual_len; i < start_addr + actual_len; i += mtd->writesize) {
        spi_message_init(&m);

        sector_write[0] = FLASH_PageProgram;
        flash_set_address(i, &sector_write[1], sector_write_size);
        t[0].tx_buf = sector_write;
        t[0].len = sector_write_size + 1;
        spi_message_add_tail(&t[0], &m);

        t[1].tx_buf = buff;
        t[1].len = rest_len > mtd->writesize ? mtd->writesize : rest_len;
        spi_message_add_tail(&t[1], &m);

        ret = spi_sync(spi, &m);
        if (ret) {
            dev_err(&spi->dev, "write sector %llu failed.ret=%d\n", i, ret);
            goto WRITE_FAILED;
        }
        rk3588_mtd_sync(mtd);
        buff += mtd->writesize;
        rest_len -= mtd->writesize;
    }

    *retlen = len;
    ret = len;
WRITE_FAILED:
    kfree(sector_write);
WRITE_SECTOR_ALLOCATE_FAILED:
    return ret;
}

int rk3588_mtd_read(struct mtd_info *mtd, loff_t from, size_t len,
                           size_t *retlen, u_char *buf)
{
    struct spi_device *spi = (struct spi_device*)mtd->priv;
    struct nm25q128evb *p = spi_get_drvdata(spi);
    u32 sector_read_size = p->deviceid == W25Q256 ? 4 : 3;
    u8 *sector_read;
    int ret;
    struct spi_message m;
    struct spi_transfer t[2];
    memset(t, 0, sizeof(struct spi_transfer) * 2);

    sector_read = kmalloc(sizeof(u8) * (sector_read_size + 2), GFP_KERNEL);
    if (!sector_read) {
        dev_err(&spi->dev, "[%s]allocate sector_read memory failed\n", __func__);
        ret = -ENOMEM;
        goto READ_SECTOR_ALLOCATE_FAILED;
    }

    spi_message_init(&m);
    sector_read[0] = FLASH_WriteDisable;
    sector_read[1] = FLASH_ReadData;
    flash_set_address(from, &sector_read[2], sector_read_size);
    t[0].tx_buf = sector_read;
    t[0].len = sector_read_size + 2;
    spi_message_add_tail(&t[0], &m);

    t[1].rx_buf = buf;
    t[1].len = len;
    spi_message_add_tail(&t[1], &m);

    ret = spi_sync(spi, &m);
    if (ret) {
        dev_err(&spi->dev, "read data failed.ret=%d\n", ret);
        goto READ_FAILED; 
    } else {
        *retlen = len;
        ret = len;
        dev_info(&spi->dev, "%s", buf);
    }
    
READ_FAILED:
    kfree(sector_read);
READ_SECTOR_ALLOCATE_FAILED:
    return ret;
}

static int parse_1_8v_power(struct spi_device *spi)
{
    struct nm25q128evb *p = spi_get_drvdata(spi);
    int ret = 0;

    ret = of_get_named_gpio(spi->dev.of_node, "gpios", 0);
    if (ret < 0) {
        dev_err(&spi->dev, "parse power gpios failed.ret=%d\n", ret);
        return ret;
    }
    p->power_gpio = ret;

    ret = gpio_request(p->power_gpio, "power");
    if (ret) {
        dev_err(&spi->dev, "request power gpio failed.ret=%d\n", ret);
        return ret;
    }

    gpio_set_value(p->power_gpio, 1);

    return 0;
}

static char const * const part_probes[] = {"ofpart"};

/*  Linux NAND Flash建立分区表的几种方式

关键词：cmdline parser、dts parser、cmdline、mtdparts等等。

MTD支持三种创建分区表的方式：

    cmdline：通过bootargs传入。
    dts：在dts中编写。
    struct mtd_partition结构体代码：代码中固定配置。

他们的优先级是：分区结构体代码 > cmdline > dts。
复制代码

mtd_device_parse_register--解析MTD分区表，并注册MTD设备。
　　parse_mtd_partitions--根据types指定的parserr名称解析分区表到parser_data中。
　　　　mtd_part_of_parse--解析DTS中分区表，使用parse_fixed_partitions()处理。。
　　　　　　mtd_part_do_parse
　　　　mtd_part_parser_get--获取parser，比如cmdline parser。
　　　　　　mtd_part_do_parse--调用parser_fn完成解析，比如parse_cmdline_partitions()。
　　　　add_mtd_partitions
　　add_mtd_partitions--在cmdline和dts partition都失败的情况下，使用代码中配置MTD分区表。 
*/
static int rk3588_spi_probe(struct spi_device *spi)
{
    int ret;
    struct nm25q128evb *p = kmalloc(sizeof(struct nm25q128evb), GFP_KERNEL);
    if (!p) {
        dev_err(&spi->dev, "allocate private data failed\n");
        goto ALLOCATE_ERR;
    }
    p->spi = spi;
    spi_set_drvdata(spi, p);

    ret = parse_1_8v_power(spi);
    if (ret) {
        goto PARSE_GPIO_ERR;
    }
    
    p->mtd.dev.parent = &spi->dev; // 要有这一句
    p->mtd.dev.of_node = spi->dev.of_node; // 要有这一句，否则mtd_device_parse_register不知道怎么解析设备树
    p->mtd._erase = rk3588_mtd_erase;
    p->mtd._write = rk3588_mtd_write;
    p->mtd._read = rk3588_mtd_read;
    p->mtd._erase = rk3588_mtd_erase;
    p->mtd._sync = rk3588_mtd_sync; // sync并没有在mtd_core中被使用
    p->mtd.type = MTD_NORFLASH;
    p->mtd.flags = MTD_CAP_NORFLASH;
    p->mtd.priv = spi;
    p->mtd.erasesize = 4096;
    p->mtd.writesize = 256;
    p->mtd.writebufsize = 256;
    p->mtd.size = 128 * 1024 * 1024; // 要设置size，mtd_add_partition的时候会对比设备树设置的偏移，超过大小会报错
    p->deviceid = flash_read_deviceid(spi);
    dev_info(&spi->dev, "spi nor flash device id = 0x%x\n", p->deviceid);

    if (p->deviceid == W25Q256) { /* SPI FLASH为W25Q256, 必须使能4字节地址模式 */
        flash_set_word(spi);
    }
    ret = mtd_device_parse_register(&p->mtd, part_probes, NULL, NULL, 0);
    if (ret) {
        dev_err(&spi->dev, "create mtd device failed.ret=%d\n", ret);
        goto ERR;
    }

    return 0;
ERR:
    gpio_free(p->power_gpio);
PARSE_GPIO_ERR:
    kfree(p);
ALLOCATE_ERR:
    return ret;
}

static int rk3588_spi_remove(struct spi_device *spi)
{
    struct nm25q128evb *p = spi_get_drvdata(spi);
    mtd_device_unregister(&p->mtd);
    gpio_free(p->power_gpio);
    return 0;
}

static const struct of_device_id g_spi_device_id[] = {
    {.compatible = "guoerba, spi-mtd-flash"},
    {},
};

static struct spi_driver g_spi_driver = {
    .probe = rk3588_spi_probe,
    .remove = rk3588_spi_remove,
    .driver = {
        .name = "mtd_spi_driver",
        .of_match_table = g_spi_device_id,
    }
};

static int __init rk3588_spi_mtd_init(void)
{
    int ret = spi_register_driver(&g_spi_driver);
    if (ret) {
        printk(KERN_ERR"spi register driver error\n");
        return ret;
    }
    return 0;
}

static void __exit rk3588_spi_mtd_exit(void)
{
    printk("unleash spi-mtd module\n");
    spi_unregister_driver(&g_spi_driver);
}

module_init(rk3588_spi_mtd_init);
module_exit(rk3588_spi_mtd_exit);
MODULE_LICENSE("GPL");
MODULE_AUTHOR("Yujie Guo");
MODULE_DESCRIPTION("rk3588-spi-mtd driver for Learning");

